import time
import torch
import os
import matplotlib.pyplot as plt
from . import model_dir, device


def plotting_results(num_epochs, train_losses, val_losses, val_accuracies):
    # 绘图
    plt.figure(figsize=(10, 6))
    plt.subplot(2, 1, 1)
    plt.plot(range(1, num_epochs + 1), train_losses, label='Train Loss')
    plt.plot(range(1, num_epochs + 1), val_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.subplot(2, 1, 2)
    plt.plot(range(1, num_epochs + 1), val_accuracies, label='Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.tight_layout()
    plt.savefig('training_metrics.png')
    # plt.show()


def train_model(train_loader, val_loader, model, criterion, optimizer, num_epochs=10, model_path='best_model.pth'):
    """
    训练和验证模型
    :param train_loader: 训练数据加载器
    :param val_loader: 验证数据加载器
    :param model: 待训练的模型
    :param criterion: 损失函数
    :param optimizer: 激活函数（优化器）
    :param num_epochs: 训练的轮数，默认为10
    :param model_path: 最佳模型保存路径
    :return:
        :param train_losses: 训练损失列表
        :param val_losses: 验证损失列表
        :param val_accuracies: 验证准确率列表
        :param train_times:
    """
    model_path = os.path.join(model_dir, model_path)  # 设定模型保存目录
    model = model.to(device)
    train_losses, val_losses, val_accuracies, train_times = [], [], [], []
    best_val_loss = float('inf')  # 初始化最佳验证损失为﹢∞

    for epoch in range(num_epochs):
        start_time = time.time()
        # 训练模式
        model.train()
        train_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()  # 梯度置零
            outputs = model(inputs)  # 正向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 计算梯度，更新权重
            train_loss += loss.item()  # 累计损失
        train_loss /= len(train_loader)  # 计算平均损失

        # 验证模式
        model.eval()
        val_loss, correct, total = 0.0, 0, 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)  # 正向传播
                loss = criterion(outputs, labels)  # 计算损失
                val_loss += loss.item()
                _, predicted = outputs.max(1)  # 计算准确率
                correct += (predicted == labels).sum().item()
                total += labels.size(0)

        val_loss /= len(val_loader)
        val_accuracy = correct / total
        train_times.append(time.time() - start_time)

        train_losses.append(train_loss)
        val_losses.append(val_loss)
        val_accuracies.append(val_accuracy)

        content = f"Epoch {epoch + 1}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}, Time: {train_times[-1]:.2f}s"
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), model_path)
            content += ", Saved Best Model!"
        else:
            content += "."
        print(content)
    return train_losses, val_losses, val_accuracies, train_times
